#include "dataset.h"
#include "bitarray.h"

void PrepareDataSet(string data_filename) {
  
  // Prepare the lexicon (and the token counts)
  Lexicon l;
  l.EasyLoad(data_filename);
  cout << "Prepared Lexicon" << endl;

  // Prepare the encoded data
  string corpus_filename = data_filename + ".corpus";
  Corpus corpus;
  if (!corpus.Load(corpus_filename)) {
    uint64 data_size = File::Size(data_filename);
    corpus.StartBuild(data_size);
    XMLTokenizer xmlt;
    xmlt.InitFromFile(data_filename);
    Token *t;
    string s;
    while ((t = xmlt.GetNextToken())) {
      s = string(t->start_, t->end_ - t->start_);
      uint32 id = l.StringToID(s.c_str());
      corpus.Push(id);
    }
    corpus.EndBuild();
    corpus.Store(corpus_filename);
  }
  corpus.DistributedTouch();
  cout << "Loaded/Prepared Corpus" << endl;
  
  // Prepare the suffixes
  string suffixes_filename = data_filename + ".suffixes";
  Memblock suffixes = Mmap2(suffixes_filename);
  if (!(suffixes.Fail())) return;
  uint64 num_tokens = corpus.size();
  IntegerArray * positions = new IntegerArray();
  positions->Init(num_tokens, CeilingLg2(num_tokens));

  for (uint64 i=0; i<num_tokens; i++) {
    if (i % 1000000 == 0) cout << "Building position list " << i << endl;
    (*positions)[i] = i;
  }
  SuffixSorter ss;
  ss.Init(&corpus, positions);
  cout << "Reasy to Sort" << endl;
  ss.Sort();
  cout << "Finished sorting, ready to store index" << endl;
  ss.Store(suffixes_filename);
  cout << "Prepared suffixes" << endl;
}

// Pass in the data filename
void DataSet::Load() {
  string dfn = ConfigString("corpus_file");

  // Load the lexicon
  CHECK(lexicon_.Load(dfn + ".lexicon"));
  // cout << "Loaded Lexicon" << endl;

  // Load the encoded data
  CHECK(corpus_.Load(dfn + ".corpus"));
  corpus_.DistributedTouch();

  // Load the suffix index
  Memblock suffixes = File::Read(dfn + ".suffixes");
  CHECK(!suffixes.Fail());
  DistributedTouch(suffixes.begin(), suffixes.size());
  index_.Init(&corpus_, suffixes);
  // cout << "Loaded suffixes" << endl;

  string token_counts_name = ConfigString("corpus_file") + ".token_counts";
  token_counts_ = File::Read(token_counts_name);
  CHECK(!token_counts_.Fail());
}
